#!/bin/bash

set -e

source "$(
  cd "$(dirname "$0")" || exit 1
  pwd
)/../common.sh"

# 安装编译前必要应用
function install_development() {
  sudo yum install -y wget tar perl-Data-Dumper autoconf automake libtool numactl binutils systemd-devel valgrind flex environment-modules gcc-c++ git
  source /etc/profile
}

# 下载HMPI
function download_HMPI() {
  echo "# Hyper MPI正在下载"
  download_HMPI_and_dependency_pkg "$1" https://github.com/kunpengcompute/hmpi/archive/refs/tags/v1.1.1-huawei.tar.gz
  download_HMPI_and_dependency_pkg "$1" https://github.com/kunpengcompute/hucx/archive/refs/tags/v1.1.1-huawei.tar.gz
  download_HMPI_and_dependency_pkg "$1" https://github.com/kunpengcompute/xucg/archive/refs/tags/v1.1.1-huawei.tar.gz
  if [[ $? -eq 0 ]]; then
    echo -e "\033[1;32;1m# Hyper MPI已经下载成功\033[0m"
  else
    echo -e "\033[31m# Hyper MPI下载失败,请重新执行脚本\033[0m"
    exit 1
  fi
}

# 下载hmpi和依赖包
function download_HMPI_and_dependency_pkg() {
  wget -t 40 -c -P "$1" "$2"
  tar xf "$1"/v1.1.1-huawei.tar.gz -C "$1"
  rm -rf "$1"/v1.1.1-huawei.tar.gz
}

# 安装hmpi
function install_HMPI() {
  echo "# Hyper MPI正在安装"
  cp -r "$1"/xucg-1.1.1-huawei/* "$1"/hucx-1.1.1-huawei/src/ucg/
  cd "$1"/hucx-1.1.1-huawei
  git init && ./autogen.sh
  ./contrib/configure-opt --prefix="$2"/hmpi-1.1.1/hucx-1.1.1-huawei --disable-numa CC=clang CXX=clang++ FC=flang
  for file in $(find . -name Makefile);do sed -i "s/-Werror//g" $file;done
  for file in $(find . -name Makefile);do sed -i "s/-implicit-function-declaration//g" $file;done
  make -j$(nproc)
  make -j$(nproc) install
  export PATH="$2"/hmpi-1.1.1/hucx-1.1.1-huawei/bin:$PATH
  export LD_LIBRARY_PATH="$2"/hmpi-1.1.1/hucx-1.1.1-huawei/lib:$LD_LIBRARY_PATH
  export INCLUDE="$2"/hmpi-1.1.1/hucx-1.1.1-huawei/include:$INCLUDE
  cd "$1"/hmpi-1.1.1-huawei
  ./autogen.pl CC=clang CXX=clang++ FC=flang
  ./configure --prefix="$2"/hmpi-1.1.1/hmpi-1.1.1-huawei --with-platform=contrib/platform/mellanox/optimized --enable-mpil-compatibility \
      --with-ucx="$2"/hmpi-1.1.1/hucx-1.1.1-huawei CC=clang CXX=clang++ FC=flang
  make -j$(nproc)
  make -j$(nproc) install
  echo -e "\033[1;32;1m# Hyper MPI已经安装完毕\033[0m"
}

# 创建HMPI的modulefiles
function write_HMPI_module() {
    cat>"$1/hmpi-1.1.1/hmpi_modulefiles"<<EOF
#%Module1.0
conflict hmpi
variable modfile [file normalize [info script]]
proc getModulefileDir {} {
    variable modfile
    set modfile_path [file dirname \$modfile]
    return \$modfile_path
}
set pwd [getModulefileDir]
set HUCX \$pwd/hucx-1.1.1-huawei
setenv HUCX \$HUCX
prepend-path PATH \$HUCX/bin
prepend-path LD_LIBRARY_PATH \$HUCX/lib
prepend-path INCLUDE \$HUCX/include
set HMPI \$pwd/hmpi-1.1.1-huawei
setenv HMPI \$HMPI
prepend-path PATH \$HMPI/bin
prepend-path LD_LIBRARY_PATH \$HMPI/lib
prepend-path INCLUDE \$HMPI/include
EOF
}

function main() {
  print_current_running_script_info "hmpi-1.1.1.sh"
  # 校验安装路径与构建路径是否一致
  is_equal_build_install_path "$1" "$2"
  # 安装编译前应用
  install_development
  # 下载
  download_HMPI "$1"
  # 安装
  install_HMPI "$1" "$2"
  # 写入modulefiles
  write_HMPI_module "$2"
  # 删除构建文件夹
  rm -rf $buildpath/xucg-1.1.1-huawei $buildpath/hucx-1.1.1-huawei $buildpath/hmpi-1.1.1-huawei 
}
main "$@"
